# -*- coding: utf-8 -*-

import os
import cv2
import math
import time
import numpy as np
from utils import stretch_linear, calc_convex_hull_2d, put_heatmap_2d, put_heatmap_3d
import warnings

# del warning
warnings.filterwarnings("ignore",message="invalid value encountered in cast")
np.seterr(divide="ignore")
def log(image):
    image_log = np.uint8(np.log(np.array(image)))
    cv2.normalize(image_log, image_log, 0, 255, cv2.NORM_MINMAX)
    # 转换成8bit图像显示
    cv2.convertScaleAbs(image_log, image_log)
    return image_log

# -------------------------------
# 计算2D凸包
img = cv2.imread("104.png", 0)
start_t = time.time()
# 均值滤波
img = cv2.blur(img, (3, 3))
# 直方图均衡化
img = cv2.equalizeHist(img)
# log
img = log(img)
# 阈值化
# _, img = cv2.threshold(img, 125, 255, cv2.THRESH_BINARY)
# 对图像进行降采样
img_arr = cv2.resize(img, (0, 0), fx=0.5, fy=0.5)
img_arr = cv2.resize(img_arr, (0, 0), fx=0.5, fy=0.5)

hull_arr = calc_convex_hull_2d(img_arr) / 255.0  # 每个scan获取2D凸包络区域
end_t = time.time()
print("Time cost(calc_convex_hull_2d and image_processing): %f s. " % (end_t - start_t))

# 计算2D凸包的轮廓
start_t = time.time()
cov_arr = hull_arr.sum() / 4.0 * np.eye(2)
y_arr, x_arr = np.where(cv2.Canny(hull_arr.astype(np.uint8) * 255, 50, 150) == 255)  # 2D凸包络区域获取轮廓像素
end_t = time.time()
print("Time cost(Canny): %f s. " % (end_t - start_t))

# 构建2D高斯热力图
start_t = time.time()
for idx in range(y_arr.shape[0]):
    mu_arr = np.array([[x_arr[idx]], [y_arr[idx]]])
    # 2D凸包络区域轮廓像素构建2D高斯热力图
    hull_arr = put_heatmap_2d(hull_arr, mu_arr, cov_arr)
end_t = time.time()
# 将降采样后的图像复原为原始尺寸
hull_arr = cv2.resize(hull_arr, (img.shape[1], img.shape[0]))
print("Time cost(put_heatmap_2d): %f s. " % (end_t - start_t))
cv2.imwrite("104_hull.png", stretch_linear(hull_arr))
# -------------------------------


scan_dir = "Pneumonia_sample/others/FE923712_les"  # mycoplasma/FE781707_les##
dst_dir = "Results1/others/FE923712_les"  # mycoplasma/FE781707_les##
scan_lst = sorted(os.listdir(scan_dir))
img_lst = []
ccinfo_dct = {"mean": [], "cov": []}
i = 0
for scan in scan_lst:
    scan_arr = cv2.imread(os.path.join(scan_dir, scan), 0)
    if (scan_arr > 127).any():
        hull_arr = calc_convex_hull_2d(scan_arr) / 255.0  # 每个scan获取2D凸包络区域
        cov_arr = hull_arr.sum() / 16.0 * np.eye(3)
        y_arr, x_arr = np.where(cv2.Canny(hull_arr.astype(np.uint8) * 255, 50, 150) == 255)  # 2D凸包络区域获取轮廓像素
        z_arr = i * np.ones_like(y_arr, dtype=np.int64)
        for j in range(z_arr.shape[0]):
            ccinfo_dct["mean"].append(np.array([[x_arr[j]], [y_arr[j]], [z_arr[j]]]))
            ccinfo_dct["cov"].append(cov_arr)
        img_lst.append(hull_arr * 255.0)
    else:
        img_lst.append(scan_arr)
    i += 1
img_arr = np.array(img_lst, dtype=np.float64) / 255.0
start_t = time.time()
for k in range(len(ccinfo_dct["mean"])):
    img_arr = put_heatmap_3d(img_arr, ccinfo_dct["mean"][k], ccinfo_dct["cov"][k])  # 2D凸包络区域轮廓像素构建3D高斯热力图
end_t = time.time()
print("Time cost: %f s. " % (end_t - start_t))

img_arr = stretch_linear(img_arr)
os.makedirs(dst_dir, exist_ok=True)
for idx in range(img_arr.shape[0]):
    cv2.imwrite(os.path.join(dst_dir, "%03d.png"%(idx+1)), img_arr[idx, :, :])
